import pandas as pd
import lightgbm as lgb
from engine.model.base_model import Model
from sklearn import preprocessing

class LGBModel(Model):
    def fit_internal(self, df, fields, train_valid_date, **kwargs):

        X_train, X_valid, y_train, y_valid  = self._prepare_data(df, fields, '2016-01-01')
        dtrain = lgb.Dataset(X_train, label=y_train)
        dvalid = lgb.Dataset(X_valid, label=y_valid)

        params = {"objective": 'mse', "verbosity": -1}
        self.model = lgb.train(
            params,
            dtrain,
            num_boost_round=1000,
            valid_sets=[dtrain, dvalid],
            valid_names=["train", "valid"],
            early_stopping_rounds=50,
            verbose_eval=True,
            #evals_result=evals_result,
            **kwargs
        )
        #evals_result["train"] = list(evals_result["train"].values())[0]
        #evals_result["valid"] = list(evals_result["valid"].values())[0]
        self.df = df

    def predict(self):
        if self.model is None:
            raise ValueError("model is not fitted yet!")
        y_pred = pd.Series(self.model.predict(self.df[self.features].values), index=self.df.index)
        # 模型评估
        from sklearn.metrics import r2_score
        print('R2可决系数',r2_score(self.df['label'], y_pred))
        return y_pred